import torch
import numpy as np


def bce(c, o):
    return np.round(-(o * np.log(c) + (1 - o) * np.log(1 - c)), 5)


if __name__ == '__main__':
    loss = torch.nn.BCELoss(reduction="none")  # reduction="none" 最后的结果不求均值
    po = [[0.1, 0.8, 0.9], [0.2, 0.7, 0.8]]
    go = [[0., 0., 1.], [0., 0., 1.]]
    p = torch.tensor(po, requires_grad=True)
    g = torch.tensor(go)
    l = loss(input=p, target=g)
    print('pytorch内置二值交叉熵损失函数：')
    print(np.round(l.detach().  # 去除梯度
                   numpy(), 5))

    pn = np.array(po)
    gn = np.array(go)
    print('自定义二值交叉熵损失函数：')
    print(bce(pn, gn))
